import os 
from operator import itemgetter
import urllib.parse


from langchain.chains.sql_database.query import create_sql_query_chain

from langchain_community.utilities import SQLDatabase
from langchain_community.tools import QuerySQLDatabaseTool
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser


os.environ['http_proxy'] = '127.0.0.1:7890'
os.environ['https_proxy'] = '127.0.0.1:7890'

os.environ["LANGSMITH_TRACING_V2"] = "true"
os.environ["LANGSMITH_API_KEY"] = "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"
os.environ["OPENAI_BASE_URL"] = "https://api.chatanywhere.tech/v1"
os.environ["OPENAI_API_KEY"] = "sk-pbXvhNj37SZ5SUBzC1Kx4LeXrsnT9EJNDL6mT2Lj2IbgohKa"
os.environ["TAVILY_API_KEY"] = "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"


model = ChatOpenAI(model='gpt-4o-mini')
# sqlalchemy 初始化MySQL数据库的连接
HOSTNAME = '127.0.0.1'
PORT = '3306'
DATABASE = 'langchain'
USERNAME = 'root'
PASSWORD = urllib.parse.quote_plus('1qaz!QAZ') 

# mysqlclient驱动URL
MYSQL_URI = 'mysql+mysqldb://{}:{}@{}:{}/{}?charset=utf8mb4'.format(USERNAME, PASSWORD, HOSTNAME, PORT, DATABASE)

db = SQLDatabase.from_uri(MYSQL_URI)
# 初始化生成SQL的chain
test_chain = create_sql_query_chain(model, db)

answer_prompt = PromptTemplate.from_template(
    """给定以下用户问题、SQL语句和SQL执行后的结果，回答用户问题。
    Question: {question}
    Query: {query}
    Result: {result}
    Answer: """
)
# 创建一个执行sql语句的工具
execute_sql_tool = QuerySQLDatabaseTool(db=db)

chain = (
    RunnablePassthrough.assign(query=test_chain)
    .assign(query=lambda x: x['query'].replace('SQLQuery:', "").strip())  # 去掉 SQLQuery: 前缀
    .assign(result=itemgetter("query") | execute_sql_tool) 
    | answer_prompt 
    | model 
    | StrOutputParser()
)
resp = chain.invoke(input={"question": "请问张飞的地址？"})
print(resp)
